import java.net.DatagramSocket;
import java.net.DatagramPacket;
import java.net.InetAddress;
import java.nio.ByteBuffer;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.TimeZone;

public class UDPClient {
    public static void main(String[] args) throws Exception {
        String serverIP = args.length > 0 ? args[0] : "localhost"; //默认localhost
        int serverPort = args.length > 1 ? Integer.parseInt(args[1]) : 9876; //默认端口9876
        DatagramSocket clientSocket = new DatagramSocket();
        InetAddress IPAddress = InetAddress.getByName(serverIP);
        clientSocket.setSoTimeout(300);//300ms超时

        //记录RTT和丢包率
        ArrayList<Long> rtts = new ArrayList<>();
        int totalSentPackets = 0;
        //数据包总数
        int totalPackets = 5;
        //每一个数据包大小
        int bytesPerPacket = 80;
        //序列号
        int sequence = 0;

        // 模拟TCP连接建立
        System.out.println("发起连接请求到 " + serverIP + ":" + serverPort + "...");
        sendPacket(clientSocket, IPAddress, serverPort, 'C', sequence++, "连接请求".getBytes("UTF-8"));
        totalSentPackets++;
        if (receiveAck(clientSocket, sequence - 1)) {
            System.out.println("连接建立成功");
        } else {
            System.out.println("连接失败");
            clientSocket.close();
            return;
        }

        //发送数据包
        for (int n = 1; n <= totalPackets; n++) {
            int startByte = (n - 1) * bytesPerPacket + 1;
            int endByte = n * bytesPerPacket;
            String content = "第 " + n + " 个（第 " + startByte + "~" + endByte + " 字节）";
            //填充数据包
            StringBuilder paddedContent = new StringBuilder(content);
            while (paddedContent.length() < bytesPerPacket) {
                paddedContent.append("0");
            }
            byte[] payload = paddedContent.toString().getBytes("UTF-8");

            //发送数据包
            boolean acknowledged = false;
            long sendTime = System.nanoTime();
            while (!acknowledged) {
                //发送数据包
                sendPacket(clientSocket, IPAddress, serverPort, 'D', sequence, payload);
                totalSentPackets++;
                //接受ACK
                acknowledged = receiveAck(clientSocket, sequence);
                if (!acknowledged) {
                    System.out.println("超时，第 " + n + " 个数据包重传 [Seq: " + sequence + "]");
                }
            }
            long receiveTime = System.nanoTime();
            //记录RTT
            long rtt = (receiveTime - sendTime) / 1000; //微秒
            rtts.add(rtt);
            System.out.println("第 " + n + " 个（第 " + startByte + "~" + endByte + " 字节）已确认, RTT: " + rtt + " μs");
            sequence++;
        }

        //计算并显示统计信息
        if (!rtts.isEmpty()) {
            long sum = 0, max = Long.MIN_VALUE, min = Long.MAX_VALUE;
            for (long rtt : rtts) {
                sum += rtt;
                max = Math.max(max, rtt);
                min = Math.min(min, rtt);
            }
            //平均RTT
            double average = (double) sum / rtts.size();
            double sumSq = 0;
            //方差
            for (long rtt : rtts) {
                sumSq += Math.pow(rtt - average, 2);
            }
            //标准差
            double stdDev = Math.sqrt(sumSq / rtts.size());
            double lossRate = ((double) (totalSentPackets - rtts.size() - 1) / totalSentPackets) * 100;

            System.out.println("\n统计信息:");
            System.out.println("丢包率: " + String.format("%.2f", lossRate) + "%");
            System.out.println("最大 RTT: " + max + " μs");
            System.out.println("最小 RTT: " + min + " μs");
            System.out.println("平均 RTT: " + String.format("%.2f", average) + " μs");
            System.out.println("RTT 标准差: " + String.format("%.2f", stdDev) + " μs");
        }

        clientSocket.close();
    }

    /**
     * 发送数据包
     * @param socket
     * @param address
     * @param port
     * @param command
     * @param sequence
     * @param payload
     * @throws Exception
     */
    private static void sendPacket(DatagramSocket socket, InetAddress address, int port, char command, int sequence, byte[] payload) throws Exception {
        ByteBuffer buffer = ByteBuffer.allocate(15 + payload.length); //首部15字节 + 数据
        buffer.put((byte) command);         //1字节命令
        buffer.putInt(sequence);            //4字节序列号
        buffer.putLong(System.currentTimeMillis()); //8字节时间戳
        buffer.putShort((short) payload.length); //2字节数据长度
        buffer.put(payload);                //数据部分
        byte[] sendData = buffer.array();
        DatagramPacket sendPacket = new DatagramPacket(sendData, sendData.length, address, port);
        socket.send(sendPacket);
    }

    /**
     * 接受ACK方法
     * @param socket
     * @param expectedSeq
     * @return
     * @throws Exception
     */
    private static boolean receiveAck(DatagramSocket socket, int expectedSeq) throws Exception {
        byte[] receiveData = new byte[1024];
        DatagramPacket receivePacket = new DatagramPacket(receiveData, receiveData.length);
        try {
            //接收数据包
            socket.receive(receivePacket);
            //解析数据包
            ByteBuffer buffer = ByteBuffer.wrap(receivePacket.getData(), 0, receivePacket.getLength());
            //命令
            char command = (char) buffer.get();
            //序列号
            int sequence = buffer.getInt();
            if (command == 'A' && sequence == expectedSeq) {
                return true;
            }
        } catch (java.net.SocketTimeoutException e) {
            return false;
        }
        return false;
    }
}